package cn.historia.test;

import com.alibaba.fastjson2.JSON;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * RAG（检索增强生成）系统测试类
 * 包含文档上传和基于检索的对话功能测试
 */
@Slf4j
@RunWith(SpringRunner.class)
@SpringBootTest
public class RAGTest {

    @Resource
    private OllamaChatClient ollamaChatClient; // Ollama大语言模型客户端
    @Resource
    private TokenTextSplitter tokenTextSplitter; // 文本分割器，用于将长文本分割成小块
    @Resource
    private SimpleVectorStore simpleVectorStore; // 简单向量存储（内存实现）
    @Resource
    private PgVectorStore pgVectorStore; // PostgreSQL向量存储（持久化实现）

    /**
     * 文档上传测试方法
     * 功能：将本地文件读取、分割并存储到PostgreSQL向量数据库中
     * 流程：
     * 1. 使用TikaDocumentReader读取指定路径下的文件
     * 2. 使用TokenTextSplitter将文档分割成适合模型处理的小块
     * 3. 为所有文档添加元数据标签，标记为"测试知识库"
     * 4. 将分割后的文档存储到PostgreSQL向量数据库
     */
    @Test
    public void upload() {
        // 创建文档读取器，指定要读取的文件路径
        TikaDocumentReader reader = new TikaDocumentReader("data/file.txt");

        // 读取文件内容，获取文档列表
        List<Document> documents = reader.get();
        // 对文档进行分割，得到适合模型处理的小块文档列表
        List<Document> documentSplitterList = tokenTextSplitter.apply(documents);

        // 为原始文档和分割后的文档添加元数据，标记为"测试知识库"
        documents.forEach(doc -> doc.getMetadata().put("knowledge", "测试知识库"));
        documentSplitterList.forEach(doc -> doc.getMetadata().put("knowledge", "测试知识库"));

        // 将分割后的文档存储到PostgreSQL向量数据库
        pgVectorStore.accept(documentSplitterList);

        // 记录上传完成日志
        log.info("上传完成");
    }


    /**
     * RAG对话测试方法
     * 功能：基于向量数据库中的知识，使用大语言模型生成回答
     * 流程：
     * 1. 定义用户查询
     * 2. 定义系统提示模板，指导模型如何使用检索到的文档
     * 3. 从向量数据库中搜索与查询相关的文档
     * 4. 构建包含检索文档的系统提示
     * 5. 调用大语言模型生成回答
     */
    @Test
    public void chat() {
        // 定义用户查询
        String message = "马也的QQ号是多少？";

        // 定义系统提示模板，指导模型如何使用检索到的文档进行回答
        // 要求模型使用文档中的信息提供准确答案，但表现得像自己本身就知道这些信息
        // 如果不确定，就简单地说不知道
        // 另外要求回复必须使用中文
        String SYSTEM_PROMPT = """
                Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
                If unsure, simply state that you don't know.
                Another thing you need to note is that your reply must be in Chinese!
                DOCUMENTS:
                    {documents}
                """;

        // 创建搜索请求，设置查询内容、返回结果数量（TopK=5）和过滤器（只搜索"测试知识库"的内容）
        SearchRequest request = SearchRequest.query(message).withTopK(5).withFilterExpression("knowledge == '测试知识库'");

        // 执行向量相似度搜索，获取与查询最相关的文档
        List<Document> documents = pgVectorStore.similaritySearch(request);
        // 将搜索到的所有文档内容拼接成一个字符串
        String documentsCollectors = documents.stream().map(Document::getContent).collect(Collectors.joining("\n"));

        // 使用系统提示模板和搜索到的文档内容创建系统消息
        Message ragMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documentsCollectors));

        // 创建消息列表，包含用户消息和系统消息
        ArrayList<Message> messages = new ArrayList<>();
        messages.add(new UserMessage(message)); // 用户消息
        messages.add(ragMessage); // 包含检索文档的系统消息

        // 调用OllamaChatClient生成回答，指定使用"deepseek-r1:1.5b"模型
        ChatResponse chatResponse = ollamaChatClient.call(new Prompt(messages, OllamaOptions.create().withModel("deepseek-r1:1.5b")));

        // 记录测试结果日志
        log.info("测试结果:{}", JSON.toJSONString(chatResponse));

    }






}
